import {Annotation, END, START, StateGraph} from "@langchain/langgraph";
import {TavilySearchResults} from "@langchain/community/tools/tavily_search";
import {getDeepseek, gpt4, qwen} from "../models";
import {HumanMessage} from "@langchain/core/messages";
import {z} from 'zod'
import { zodToJsonSchema } from 'zod-to-json-schema'
import {ChatOpenAI} from "@langchain/openai";
import {createReactAgent} from "@langchain/langgraph/prebuilt";
import {ChatPromptTemplate} from "@langchain/core/prompts";
import {JsonOutputToolsParser} from "@langchain/core/output_parsers/openai_tools";
import {RunnableConfig} from "@langchain/core/runnables";
import {saveImages} from "../tool";

const deepseek = getDeepseek()

const PlanExecuteState = Annotation.Root({
    input: Annotation<string>({
        reducer: (x, y) => y ?? x ?? "",
    }),
    plan: Annotation<string[]>({
        reducer: (x, y) => y ?? x ?? [],
    }),
    pastSteps: Annotation<[string, string][]>({
        reducer: (x, y) => x.concat(y),
    }),
    response: Annotation<string>({
        reducer: (x, y) => y ?? x,
    }),
})

const tools = [new TavilySearchResults({ maxResults: 3 ,apiKey: "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"})];

const agentExecutor = createReactAgent({
    llm: qwen,
    tools: tools,
});

// const rep = await agentExecutor.invoke({
//     messages: [new HumanMessage("who is the winner of the us open")]
// });
// console.log(rep)


const plan = zodToJsonSchema(
    z.object({
        steps: z
            .array(z.string())
            .describe("different steps to follow, should be in sorted order"),
    }),
);
const planFunction = {
    name: "plan",
    description: "This tool is used to plan the steps to follow",
    parameters: plan,
};

const planTool = {
    type : 'function',
    function: planFunction
}

const plannerPrompt = ChatPromptTemplate.fromTemplate(
    `For the given objective, come up with a simple step by step plan. \
This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \
The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.

{objective}
`
)

const model = qwen.withStructuredOutput(plan)

const planner = plannerPrompt.pipe(model)
//
// const planRep= await planner.invoke({
//     objective: "what is the hometown of the current Australia open winner"
// })
//
// console.log(planRep);


const response = zodToJsonSchema(
    z.object({
        response: z.string().describe('Response to user')
    })
)

const responseTool = {
    type: 'function',
    function: {
        name: 'response',
        description: 'Response to user.',
        parameters: response
    }
}

const replannerPrompt = ChatPromptTemplate.fromTemplate(
    `For the given objective, come up with a simple step by step plan. 
This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps.
The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.

Your objective was this:
{input}

Your original plan was this:
{plan}

You have currently done the follow steps:
{pastSteps}

Update your plan accordingly. If no more steps are needed and you can return to the user, then respond with that and use the 'response' function.
Otherwise, fill out the plan.  
Only add steps to the plan that still NEED to be done. Do not return previously done steps as part of the plan.\`,

`,
);

const parser = new JsonOutputToolsParser();

const replanner = replannerPrompt.pipe(
    qwen.bindTools([
        responseTool,
        planTool
    ])
).pipe(parser)

async function executeStep(state: typeof PlanExecuteState.State, config?: RunnableConfig):Promise<Partial<typeof PlanExecuteState.State>> {
    // console.log('plan:',state.plan)
    const task = state.plan[0];
    const input = {
        messages: [new HumanMessage(task)]
    }
    const { messages } = await agentExecutor.invoke(input, config)

    return {
        pastSteps: [
            [task, messages[messages.length -1].content.toString() ],
        ],
        plan:state.plan.slice(1)
    }
}

async function planStep(state: typeof PlanExecuteState.State):Promise<Partial<typeof PlanExecuteState.State>> {
    const plan = await planner.invoke({objective: state.input})
    return  {
        plan: plan.steps
    }
}

async function replanStep(state: typeof PlanExecuteState.State):Promise<Partial<typeof PlanExecuteState.State>>  {
    const output = await replanner.invoke({
        input: state.input,
        plan: state.plan.join('\n'),
        pastSteps: state.pastSteps.map(([step, result])=> `${step}: ${result}`).join('\n')
    })

    const toolCall = output[0]
    if(toolCall.type == 'response') {
        return {
            response: toolCall.args?.response
        }
    }
    return {
        plan: toolCall.args?.steps
    }
}

function shouldEnd(state: typeof PlanExecuteState.State) {
    return state.response? 'true': 'false'
}

const workflow = new StateGraph(PlanExecuteState)
.addNode('planner', planStep)
.addNode('agent', executeStep)
.addNode('replan', replanStep)
.addEdge(START, 'planner')
.addEdge('planner', 'agent')
.addEdge('agent', 'replan')
.addConditionalEdges('replan', shouldEnd, {
    true: END,
    false: 'agent'
})

const app = workflow.compile()

saveImages(app, "plan.png")

// const config = { recursionLimit: 50 };
// const inputs = {
//     input: "蔚来最近4年营收对比?",
// };
//
// for await (const event of await app.stream(inputs, config)) {
//     console.log(event);
// }

//
// const pastSteps = [
//  [ "Identify the winner of the 2024 Australian Open", "The winner of the 2024 Australian Open men's singles title is Jannik Sinner from Italy. He defeated Daniil Medvedev in the final match." ],
// [ "Find the hometown of Jannik Sinner", "Based on the search results, Jannik Sinner's hometown is Sesto Pusteria, which is located in the county of Trentino-Alto Adige in northern Italy, near the Alps. This information can be found in both articles mentioning his return to his hometown after reaching the top ranking in tennis." ]
// ]
// const plans = ['Identify the winner of the 2024 Australian Open',"Find the hometown of Jannik Sinner"]
// const test = await replanner.invoke({
//     input: "what is the hometown of the 2024 Australian open winner?",
//     plan: "",
//     pastSteps: pastSteps.map(([step, result])=> `${step}: ${result}`).join('\n')
// })
//
// console.log(test)